from transformers import BertTokenizer, BertForMaskedLM, pipeline

# 加载 tokenizer 和模型
tokenizer = BertTokenizer.from_pretrained('D:\\work\program\\pytorch_models\\bert-base-chinese')
model = BertForMaskedLM.from_pretrained('D:\\work\program\\pytorch_models\\bert-base-chinese')

# 创建文本生成管道
nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)

# 输入带掩码的文本
text = "北京是[MASK]的首都。"

# 进行文本生成
result = nlp(text)
print(result)